# 导入必要的库
import numpy as np
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from sklearn.preprocessing import StandardScaler

# 数据预处理函数
def preprocess_data(file_path):
    # 加载数据
    data = np.loadtxt(file_path)
    # 数据分割
    X = data[:, :16]
    y = data[:, 16:]
    # 数据标准化
    scaler_X = StandardScaler()
    X_scaled = scaler_X.fit_transform(X)
    return X_scaled, y, scaler_X

# 构建神经网络模型函数
def build_model():
    model = Sequential([
        Dense(128, input_dim=16, activation='relu'),
        Dense(64, activation='relu'),
        Dense(6, activation='linear')
    ])
    model.compile(loss='mean_squared_error', optimizer='adam')
    return model

# 主函数
def main():
    # 数据预处理
    X_train_scaled, y_train, scaler_X = preprocess_data('./data/train_data.txt')
    X_test_scaled, y_test, _ = preprocess_data('./data/test_data.txt')
    
    # 构建和训练模型
    model = build_model()
    model.fit(X_train_scaled, y_train, epochs=50, batch_size=128, validation_split=0.2)
    
    # 预测测试集
    y_pred_scaled = model.predict(X_test_scaled)
    
    # 反标准化预测值（这里不需要，因为我们的目标是MSE，而MSE是在原始尺度上计算的）
    # y_pred = scaler_y.inverse_transform(y_pred_scaled)
    
    # 计算MSE
    mse = np.mean((y_pred_scaled - y_test) ** 2)
    adjusted_mse = mse / 2
    
    # 打印和保存MSE
    print(f'Mean Squared Error on test data: {mse}')
    print(f'Adjusted Mean Squared Error on test data: {adjusted_mse}')
    
# 以下是保存到单独文件中
# with open('mse_result.txt', 'w') as f:
# f.write(f'Mean Squared Error on test data: {mse}\n')
# f.write(f'Adjusted Mean Squared Error on test data: {adjusted_mse}\n')

    # 将MSE保存到test_data.txt文件的末尾
    with open('./data/test_data.txt', 'a') as f:
        f.write(f'\nMean Squared Error on test data: {mse}')

# 执行主函数
if __name__ == '__main__':
    main()
